import { ChatMessage, DiffLine } from "..";
import { filterLeadingAndTrailingNewLineInsertion, filterEnglishLinesAtStart, filterCodeBlockLines, stopAtLines, skipLines, removeTrailingWhitespace } from "../autocomplete/filtering/streamTransformes/lineStream";
import { streamLines } from "../diff/util";
import { streamDiff } from "../diff/streamDiff";

export type LineStream = AsyncGenerator<string>;


export async function* streamDiffLines(
  prefix: string,
  highlighted: string,
  completion: AsyncGenerator<ChatMessage>
): AsyncGenerator<DiffLine> {

  // Strip common indentation for the LLM, then add back after generation
  let oldLines =
    highlighted.length > 0
      ? highlighted.split("\n")
      :[];

  // But if that line is empty, we can assume we are insertion-only
  if (oldLines.length === 1 && oldLines[0].trim() === "") {
    oldLines = [];
  }

  // Trim end of oldLines, otherwise we have trailing \r on every line for CRLF files
  oldLines = oldLines.map((line) => line.trimEnd());


  let lines = streamLines(completion);

  lines = filterEnglishLinesAtStart(lines);
  lines = filterCodeBlockLines(lines);
  lines = stopAtLines(lines, () => {});
  lines = skipLines(lines);
  lines = removeTrailingWhitespace(lines);
  // if (inept) {
  //   // lines = fixCodeLlamaFirstLineIndentation(lines);
  //   lines = filterEnglishLinesAtEnd(lines);
  // }

  let diffLines = streamDiff(oldLines, lines);
  diffLines = filterLeadingAndTrailingNewLineInsertion(diffLines);

  const line = highlighted.split("\n").filter(line => line !== "").slice(0, 1)[0];
  const indentation = line.slice(0, line.length - line.trimStart().length);
  diffLines = addIndentation(diffLines, indentation);

  let onlyOneInsertion = false;
  let seenGreen = false;
  for await (const diffLine of diffLines) {
    yield diffLine;
    if (diffLine.type === "new") {
      seenGreen = true;
    } else if (onlyOneInsertion && seenGreen && diffLine.type === "same") {
      break;
    }
  }
}

export async function* addIndentation(
  diffLineGenerator: AsyncGenerator<DiffLine>,
  indentation: string,
): AsyncGenerator<DiffLine> {
  for await (const diffLine of diffLineGenerator) {
    if (diffLine.type == "new") {
      yield {
        ...diffLine,
        line: indentation + diffLine.line,
      };
    } else {
      yield diffLine;
    }
  }
}